import os
import argparse
import sys
from imagededup.methods import PHash, DHash, CNN
import shutil


parser = argparse.ArgumentParser("video2images")

parser.add_argument("--video_path", '-v', type=str, help='The video path.')
parser.add_argument("--save_dir", '-s', type=str, default='./output', help='The save dir to images')
parser.add_argument("--vf", type=int, default=10, help='save picture in one second')


def use_ffmpeg(_video_path, _save_dir, vf=10):
    abs_dir = os.path.abspath(sys.path[0])
    cmd = abs_dir + '\\ffmpeg\\bin\\ffmpeg.exe -i ' + _video_path + ' -vf fps=1/' + str(vf) + ' -f image2 ' + _save_dir + '%3d.jpg'
    print(cmd)
    os.system(cmd)


def deduplicate(_save_dir):
    image_dir = _save_dir
    phasher = CNN()

    # 生成图像目录中所有图像的二值hash编码
    encodings = phasher.encode_images(image_dir=image_dir)

    # find_duplicates方法 对已编码图像寻找重复图像
    # duplicates = phasher.find_duplicates(encoding_map=encodings)
    duplicates = phasher.find_duplicates_to_remove(encoding_map=encodings)

    # plot_duplicates(image_dir=image_dir,
    #                 duplicate_map=duplicates,
    #                 filename='098.jpg')

    print("[INFO] 将要删除以下图像")
    print(duplicates)

    os.mkdir(image_dir + 'res')
    file_or_dir = os.listdir(image_dir)
    for file_dir in file_or_dir:
        if not os.path.isdir(file_dir) and file_dir not in duplicates:
            shutil.copyfile(image_dir + file_dir, image_dir + 'res/' + file_dir)

    # for each_file in duplicates:
    #     os.remove(image_dir + each_file)


if __name__ == '__main__':
    args = parser.parse_args()
    video_path = args.video_path
    save_path = args.save_dir

    if not video_path:
        print("video_path null")

    if not save_path:
        print("save_path null")

    if save_path[-1] != '/' or save_path[-1] != '\\':
        if '/' in save_path:
            save_path += '/'
        else:
            save_path += '\\'

    print("[INFO] 视频取帧开始")
    use_ffmpeg(video_path, save_path, args.vf)

    print("[INFO] 取帧去重开始")
    deduplicate(save_path)

    print('[INFO] OCR识别开始')
    # use_ocr(save_path)

    print("[INFO] 运行结束")
